#include "generate_code.h"


const uint32_t op_mov_imm32[REG_MAX] = {
    [RAX] = 0xc0c748, [RBX] = 0xc3c748, [RCX] = 0xc1c748, [RDX] = 0xc2c7c8, [RDI] = 0xc7c748,
    [R8] = 0xc0c7c9,  [R9] = 0xc1c7c9,  [R10] = 0xc2c7c9, [R11] = 0xc3c7c9, [R12] = 0xc4c7c9,
    [R13] = 0xc5c7c9, [R14] = 0xc6c7c9, [R15] = 0xc7c7c9
};

const uint32_t op_movabs_imm64[REG_MAX] = {
    [RAX] = 0xb848, [RBX] = 0xbb48, [RCX] = 0xb948, [R13] = 0xbd49, [R14] = 0xbe49, [R15] = 0xbf49
};

const uint32_t op_dec_reg[REG_MAX] = {
    [RAX] = 0xc8ff48, [RBX] = 0xcbff48, [RCX] = 0xc0ff48, [RDX] = 0xcaff48, [RDI] = 0xcfff48
};

size_t mov_reg_imm32(uint8_t *dst, REG_NO reg, uint32_t imm32)
{
    *(uint32_t *)dst = op_mov_imm32[reg];
    *(uint32_t *)&dst[3] = imm32;
    return 7;
}

size_t mov_reg_imm64(uint8_t *dst, REG_NO reg, uint64_t imm64)
{
    *(uint16_t *)dst = op_movabs_imm64[reg];
    *(uint64_t *)&dst[2] = imm64;
    return 10;
}

size_t dec_reg(uint8_t *dst, REG_NO reg)
{
    *(uint32_t *)dst = op_dec_reg[reg];
    return 3;
}

size_t jnz_imm32(uint8_t *dst, uint32_t offset)
{
    *(uint16_t *)dst = 0x850f;
    *(uint32_t *)&dst[2] = offset - 6;
    return 6;
}

size_t shl_rdx_imm8(uint8_t *dst, uint8_t offset)
{
    // shl rdx, imm8
    if (offset == 1) {
        *(uint32_t *)dst = 0xe2d148;
        return 3;
    }
    *(uint32_t *)dst = offset << 24 | 0xe2c148;
    return 4;
}

size_t lfence(uint8_t *dst)
{
    *(uint32_t *)dst = 0xe8ae0f;
    return 3;
}

size_t rdtsc(uint8_t *dst)
{
    *(uint16_t *)dst = 0x310f;
    return 2;
}

size_t rdpmc(uint8_t *dst)
{
    *(uint16_t *)dst = 0x330f;
    return 2;
}

size_t splice_rdx_rax(uint8_t *dst)
{
    // rdx << 32 | rax
    size_t len = shl_rdx_imm8(dst, 32);
    *(uint32_t *)&dst[len] = 0xc20948;
    return len + 3;
}

size_t sub_rdx_rax_to_r15(uint8_t *dst, uint8_t offset)
{
    // sub (rdx << 32 | rax), [r15]
    size_t len = splice_rdx_rax(dst);
    if (offset == 0) {
        *(uint32_t *)&dst[len] = 0x172949;
        return len + 3;
    }
    *(uint32_t *)&dst[len] = offset << 24 | 0x572949;
    return len + 4;
}

size_t add_rdx_rax_to_r15(uint8_t *dst, uint8_t offset)
{
    // add (rdx << 32 | rax), [r15]
    size_t len = splice_rdx_rax(dst);
    if (offset == 0) {
        *(uint32_t *)&dst[len] = 0x170149;
        return len + 3;
    }
    *(uint32_t *)&dst[len] = offset << 24 | 0x570149;
    return len + 4;
}

size_t reg_push(uint8_t *dst)
{
    // r15 to restore result
    // rdi to count loop

    // push rbp
    *dst = 0x55;
    dst += 1;

    // mov rbp, rsp
    *(uint32_t *)dst = 0xe58948;
    dst += 3;

    // push r15
    *(uint16_t *)dst = 0x5741;
    dst += 2;

    // push rdi
    *(uint8_t *)dst = 0x57;

    return 7;
}

size_t reg_pop(uint8_t *dst)
{
    // pop rdi
    *(uint8_t *)dst = 0x5f;
    dst += 1;

    // pop r15
    *(uint16_t *)dst = 0x5f41;
    dst += 2;

    // pop rbp
    *dst = 0x5d;
    return 4;
}
